Amazon SageMaker HyperPod で PyTorch による分散学習を試してみた

CLUSTER_ID=jt52saf5imim
NODE_GROUP_NAME=controller-machine
INSTANCE_ID=i-086f3d3670a48a658

aws ssm start-session --target sagemaker-cluster:${CLUSTER_ID}_${NODE_GROUP_NAME}-${INSTANCE_ID}

$ sinfo
PARTITION     AVAIL  TIMELIMIT  NODES  STATE NODELIST
dev*             up   infinite      2   idle ip-10-1-86-158,ip-10-1-120-53
ml.c5.4xlarge    up   infinite      2   idle ip-10-1-86-158,ip-10-1-120-53

#!/bin/bash
#SBATCH --job-name=cpu-ddp-conda  # ジョブ名を付ける
#SBATCH --exclusive               # 他のジョブとノードを共有しない
#SBATCH --wait-all-nodes=1       # 1 を指定するとすぐに実行せずに全ノードが利用可能になるまで待機
#SBATCH --nodes 2                # 使用するノード数を指定
#SBATCH --output=logs/%x_%j.out  # ログ出力設定（出力形式: [sbatch を実行時のカレントリディレクトリ]/logs/[ジョブ名]_[ジョブ ID].out）

declare -a TORCHRUN_ARGS=(
    --nproc_per_node=4           # ノードごとのプロセス数
    --nnodes=$SLURM_JOB_NUM_NODES # Slurmの変数から取得した総ノード数を指定(今回は 2 ノード)
    --rdzv_backend=c10d          # ランデブーを処理するために使用するバックエンドの指定
    --rdzv_endpoint=$(hostname)   # ランデブーのエンドポイント（ホスト名）の指定
)

if [ -d "/opt/sagemaker_cluster" ]; then
    echo "Detected Hyperpod cluster.. enabling --auto-resume=1"
    AUTO_RESUME="--auto-resume=1"
fi

srun ${AUTO_RESUME} ./pt_cpu/bin/torchrun \
    "${TORCHRUN_ARGS[@]}" \
    $(dirname "$PWD")/ddp.py 5000000 10

# モデルの定義部分
model = torch.nn.Linear(20, 1)  # 20個の値の入力から1つの値を出力する線形モデル

# 入力データの生成
self.data = [(torch.rand(20), torch.rand(1)) for _ in range(size)]

def ddp_setup():
    init_process_group(backend="gloo") # GPU 用は "nccl" を指定

srun ${AUTO_RESUME} ./pt_cpu/bin/torchrun \
    "${TORCHRUN_ARGS[@]}" \
    $(dirname "$PWD")/ddp.py 5000000 10

$ sbatch 1.conda-train.sbatch
Submitted batch job 1

$ squeue
             JOBID PARTITION     NAME     USER ST       TIME  NODES NODELIST
                 1       dev cpu-ddp-   ubuntu  R       0:01      2 ip-10-1-86-158,ip-10-1-120-53

CLUSTER_ID=jt52saf5imim
NODE_GROUP_NAME=worker-group-1
INSTANCE_ID=i-0fd1c39229f2a1e0f

aws ssm start-session --target sagemaker-cluster:${CLUSTER_ID}_${NODE_GROUP_NAME}-${INSTANCE_ID}

$ squeue
             JOBID PARTITION     NAME     USER ST       TIME  NODES NODELIST(REASON)
                11       dev cpu-ddp-   ubuntu  R 1-12:40:42      2 ip-10-1-86-158,ip-10-1-120-53

[RANK 4] Epoch 4999999 | Batchsize: 32 | Steps: 8[RANK 6] Epoch 4999999 | Batchsize: 32 | Steps: 8

[RANK 7] Epoch 4999999 | Batchsize: 32 | Steps: 8
[RANK 3] Epoch 4999999 | Batchsize: 32 | Steps: 8
INFO:torch.distributed.elastic.agent.server.api:[default] worker group successfully finished. Waiting 300 seconds for other agents to finish.
INFO:torch.distributed.elastic.agent.server.api:Local worker group finished (SUCCEEDED). Waiting 300 seconds for other agents to finish
INFO:torch.distributed.elastic.agent.server.api:[default] worker group successfully finished. Waiting 300 seconds for other agents to finish.
INFO:torch.distributed.elastic.agent.server.api:Local worker group finished (SUCCEEDED). Waiting 300 seconds for other agents to finish
INFO:torch.distributed.elastic.agent.server.api:Done waiting for other agents. Elapsed: 0.0011348724365234375 seconds
INFO:torch.distributed.elastic.agent.server.api:Done waiting for other agents. Elapsed: 4.624900579452515 seconds

$ ls -lh snapshot.pt
-rw-rw-r-- 1 ubuntu ubuntu 1.2K Dec 30 20:16 snapshot.pt